package com.windmt.filter;

import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Bucket4j;
import io.github.bucket4j.Refill;
import lombok.Data;
import lombok.Getter;
import lombok.Setter;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.core.Ordered;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;

import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author Jean
 * @date
 * @des 使用Bucket4J 实现限流/SpringCloud 也有自带的限流过滤器,详见(RequestRateLimiterGatewayFilterFactory)
 */
@Component
public class RateLimitByIpGatewayFilterFactory extends AbstractGatewayFilterFactory<RateLimitByIpGatewayFilterFactory.Config> implements Ordered {

    private final Log log = LogFactory.getLog(this.getClass());
    /**
     * 令牌桶容量
     */
    private static int burstCapacity;

    /**
     * 令牌每次补充数量
     */
    private static int replenishRate;

    /**
     * 令牌补充间隔
     */
    private static Duration refillDuration;

    /**
     * 令牌桶放到内存
     */
    private static final Map<String, Bucket> CACHE = new ConcurrentHashMap<>();

    /**
     * 新增令牌桶
     *
     * @return
     */
    private Bucket createNewBucket() {
        Refill refill = Refill.of(replenishRate, refillDuration);
        Bandwidth limit = Bandwidth.classic(burstCapacity, refill);
        return Bucket4j.builder().addLimit(limit).build();
    }


    @Override
    public GatewayFilter apply(Config config) {
        return (exchange, chain) -> {
            String hostAddress = exchange.getRequest().getRemoteAddress().getAddress().getHostAddress();
            Bucket bucket = CACHE.computeIfAbsent(hostAddress, k -> createNewBucket());

            log.debug("IP: " + hostAddress + "，TokenBucket Available Tokens: " + bucket.getAvailableTokens());

            if (bucket.tryConsume(1)) {
                return chain.filter(exchange);
            }
            ServerHttpResponse response = exchange.getResponse();
            response.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
            return response.setComplete();
        };
    }

    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE;
    }

    public RateLimitByIpGatewayFilterFactory() {
        super(Config.class);
    }

    @Getter
    public static class Config {
        private int burstCapacity;
        private int replenishRate;
        private Duration refillDuration;

        public Config setBurstCapacity(int burstCapacity) {
            this.burstCapacity = burstCapacity;
            RateLimitByIpGatewayFilterFactory.burstCapacity = burstCapacity;
            return this;
        }

        public Config setReplenishRate(int replenishRate) {
            this.replenishRate = replenishRate;
            RateLimitByIpGatewayFilterFactory.replenishRate = replenishRate;
            return this;
        }

        public Config setRefillDuration(Duration refillDuration) {
            this.refillDuration = refillDuration;
            RateLimitByIpGatewayFilterFactory.refillDuration = refillDuration;
            return this;
        }
    }
}
